import torch
import os

def auto_device():
    if not torch.cuda.is_available():
        return "cpu"
    # 获取当前设备的数量
    device_count = torch.cuda.device_count()
    print("device_count:", device_count)
    
    most_free_device = -1
    most_free_memory = -1
    for i in range(device_count):
        # 获取当前设备的空余显存
        memory_info = torch.cuda.get_device_properties(i).total_memory - torch.cuda.memory_reserved(i) - torch.cuda.memory_allocated(i)
        print(f"device {i} total_memory:", torch.cuda.get_device_properties(i).total_memory)
        print(f"device {i} memory_reserved:", torch.cuda.memory_reserved(i))
        print(f"device {i} memory_allocated:", torch.cuda.memory_allocated(i))
        print(f"device {i} memory_info:", memory_info)
        if memory_info > most_free_memory:
            most_free_memory = memory_info
            most_free_device = i
    print("most_free_device:", most_free_device)
    return f"cuda:{most_free_device}"



if __name__ == "__main__":
    auto_device()